#include "../include/senyun_driver/undistort.h"
#include "cuda_runtime.h"
#include <string>
#include <iostream>

using namespace std;

__global__ void undistort_kernel(unsigned char* src_img, unsigned char* out_img, float k1, float k2, float p1, float p2, float k3, float fx, float fy, float cx, float cy, int cols, int rows)
{
    int u_distorted = 0, v_distorted = 0;
    float x1,y1,x2,y2,r2;
    int index = threadIdx.x + blockIdx.x * blockDim.x;
    if (index>=rows*cols)
    {
	return;
    } 
    int undistort_index= 0;
    int output_index = 0;
    //get row index and col index
    int u = index%cols;
    int v = index/cols;
    //
    x1 = (u-cx)/fx;
    y1 = (v-cy)/fy;
    r2 = x1*x1+y1*y1;
    float k3r3 = 1 + k1 * r2 + k2 * r2*r2+k3 * r2*r2*r2;
    
    x2 = x1*k3r3+ 2*p1*x1*y1 + p2 * (r2+2*x1*x1);
    y2 = y1*k3r3 +  p1*(r2+2*y1*y1)+2*p2*x1*y1;
    u_distorted = fx*x2+cx;
    v_distorted = fy*y2+cy;
    output_index = index*3;
    //if (u_distorted >= 0 && v_distorted >= 0 && u_distorted < cols && v_distorted < rows) 
    //{
    undistort_index = (v_distorted)*cols*3 + (u_distorted)*3;
    out_img[output_index] = src_img[undistort_index];
    out_img[output_index+1] = src_img[undistort_index+1];
    out_img[output_index+2] = src_img[undistort_index+2];
	
    //} 
    //else
    //{
//	    out_img[output_index] = 0;
//	    out_img[output_index+1] = 0;
//	    out_img[output_index+2] = 0;
//	}
      
}

void undiundistort(float* param, unsigned char* input_image_gpu, unsigned char* output_image_gpu, int rows, int cols)
{
    float k1 = param[0];
    float k2 = param[1];
    float p1 = param[2];
    float p2 = param[3];
    float k3 = param[4];
    float fx = param[5];
    float fy = param[6];
    float cx = param[7];
    float cy = param[8];
    //float k3r3 = 1 + k1 * r2 + k2 * r2*r2+k3 * r2*r2*r2;;
    int dim_grid=(rows*cols+1024-1)/1024;
    undistort_kernel<<<dim_grid,1024>>>(input_image_gpu, output_image_gpu,k1,k2,p1,p2,k3,fx,fy,cx,cy,cols,rows);
    //cudaDeviceSynchronize();

}



